import tensorboardX
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from PIL import Image
import torch
from models import Discriminator, Generator
from utils import LambdaLR, ReplayBuffer, weights_init_normal
from image_datasets import ImageDataset
import itertools
from torchvision.utils import save_image

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

batch_size = 1
size = 256
lr = 0.0002
n_epoch = 200
epoch = 0
decay_epoch = 100
step = 0

# networks
netG_A2B = Generator().to(device)
netG_A2B.load_state_dict(torch.load("save_models/netG_A2B.pkl"))
netG_A2B.eval()

netG_B2A = Generator().to(device)
netG_B2A.load_state_dict(torch.load("save_models/netG_B2A.pkl"))
netG_B2A.eval()

input_A = torch.ones([1, 3, size, size], dtype=torch.float).to(device)
input_B = torch.ones([1, 3, size, size], dtype=torch.float).to(device)

transforms_ = [
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]

data_root = "apple2orange"
data_loader = DataLoader(ImageDataset(data_root, transforms_, "test"), batch_size=batch_size, shuffle=False,
                         num_workers=0)

for i, batch in enumerate(data_loader):
    real_A = torch.tensor(input_A.copy_(batch['A']), dtype=torch.float).to(device)
    real_B = torch.tensor(input_B.copy_(batch['B']), dtype=torch.float).to(device)

    fake_B = 0.5 * (netG_A2B(real_A).data + 1.0)
    fake_A = 0.5 * (netG_B2A(real_B).data + 1.0)

    save_image(fake_A, "outputs/A/{}.png".format(i))
    save_image(fake_B, "outputs/B/{}.png".format(i))

    print(i)
